package com.dunshan.dunshanmongodb.config;

import com.alibaba.ttl.TransmittableThreadLocal;
import com.mongodb.client.MongoDatabase;
import lombok.extern.slf4j.Slf4j;
import org.springframework.data.mongodb.MongoDbFactory;
import org.springframework.data.mongodb.core.MongoTemplate;

@Slf4j
public class MultiMongoTemplate extends MongoTemplate {
    //    用来缓存当前MongoDbFactory
    static TransmittableThreadLocal<MongoDbFactory> mongoDbFactoryThreadLocal = new TransmittableThreadLocal<>();

    //用来缓存当前MongoDbFactory
//    private static ThreadLocal<MongoDbFactory> mongoDbFactoryThreadLocal;

    public MultiMongoTemplate(MongoDbFactory mongoDbFactory) {
        super(mongoDbFactory);
        if (mongoDbFactoryThreadLocal == null) {
            mongoDbFactoryThreadLocal = new TransmittableThreadLocal<>();
        }
    }

    public void setMongoDbFactory(MongoDbFactory factory) {
        mongoDbFactoryThreadLocal.set(factory);
    }

    public void removeMongoDbFactory() {
        mongoDbFactoryThreadLocal.remove();
    }

    @Override
    public MongoDatabase getDb() {
        log.info("目前数据库是：{}", mongoDbFactoryThreadLocal.get().getDb());
        return mongoDbFactoryThreadLocal.get().getDb();
    }
}

